from ultralytics import YOLO
from argparse import ArgumentParser

def parse_args():
    parser = ArgumentParser()
    parser.add_argument('--model', default='yolo11n.pt')
    parser.add_argument('--imgsz', default=640, type=int)
    parser.add_argument('--device', default='0')
    parser.add_argument('--batch', default=16, type=int)
    parser.add_argument('--workers', default=0, type=int)
    parser.add_argument('--memory-cache', action='store_true', )
    parser.add_argument('--data', default='uav_iscas.yaml')
    return parser.parse_args()
def main():
    #    Load a model
    args = parse_args()
    model = YOLO(model=args.model)

    # Train the model
    train_results = model.train(
        data=args.data,  # path to dataset YAML
        epochs=100,  # number of training epochs
        patience=10,
        imgsz=args.imgsz,  # training image size
        device=args.device,  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu,
        batch=args.batch,  # -1 for auto batch size
        cache=args.memory_cache,
        workers=args.workers
    )

    # Evaluate model performance on the validation set
    # val_metrics = model.val()
    test_metrics = model.val(split='test')

    # print(val_metrics)
    # print(test_metrics)

if __name__ == '__main__':
    main()